"""
CUDA 检测脚本
检查系统 CUDA 安装情况和显卡信息
"""

import subprocess
import sys
import os

print("=" * 70)
print("CUDA 环境检测工具")
print("=" * 70)
print()

# 检查 NVIDIA 驱动
print("[步骤 1/4] 检查 NVIDIA 驱动...")
try:
    result = subprocess.run(
        ["nvidia-smi"],
        capture_output=True,
        text=True,
        timeout=10
    )
    
    if result.returncode == 0:
        print("✓ NVIDIA 驱动已安装")
        print()
        print(result.stdout)
        
        # 尝试提取 CUDA 版本
        if "CUDA Version:" in result.stdout:
            cuda_version = result.stdout.split("CUDA Version:")[1].split()[0]
            print(f"✓ 检测到 CUDA 版本: {cuda_version}")
            print()
    else:
        print("✗ 未检测到 NVIDIA 驱动")
        print()
        print("请先安装 NVIDIA 显卡驱动")
        print("下载地址: https://www.nvidia.com/Download/index.aspx")
        sys.exit(1)
        
except FileNotFoundError:
    print("✗ 未找到 nvidia-smi 命令")
    print()
    print("可能的原因:")
    print("1. NVIDIA 驱动未安装")
    print("2. 驱动安装不完整")
    print()
    print("解决方法:")
    print("1. 访问 https://www.nvidia.com/Download/index.aspx")
    print("2. 选择您的显卡型号（GTX 1660）")
    print("3. 下载并安装最新驱动")
    sys.exit(1)
    
except Exception as e:
    print(f"✗ 检测失败: {str(e)}")
    sys.exit(1)

# 检查当前 PyTorch
print("[步骤 2/4] 检查当前 PyTorch...")
try:
    import torch
    print(f"✓ PyTorch 已安装 (版本: {torch.__version__})")
    
    if torch.cuda.is_available():
        print(f"✓ PyTorch 已支持 CUDA")
        print(f"  CUDA 版本: {torch.version.cuda}")
        print(f"  可用 GPU 数量: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
        print()
        print("=" * 70)
        print("✓✓✓ 您的 PyTorch 已经支持 GPU！✓✓✓")
        print("=" * 70)
        print()
        print("无需重新安装，可以直接使用。")
        sys.exit(0)
    else:
        print(f"⚠ PyTorch 不支持 CUDA (当前是 CPU 版本)")
        print(f"  需要卸载并重新安装 GPU 版本")
        
except ImportError:
    print("⚠ PyTorch 未安装")
    
print()

# 推荐安装命令
print("[步骤 3/4] 分析系统配置...")
print()
print("您的显卡: NVIDIA GTX 1660")
print("计算能力: 7.5 (支持 CUDA)")
print()

# 检测 Python 版本
python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
print(f"Python 版本: {python_version}")
print()

print("[步骤 4/4] 推荐安装方案")
print("=" * 70)
print()

# 推荐 CUDA 11.8 (兼容性好，稳定)
print("【推荐方案】CUDA 11.8 版本（兼容性最好）")
print()
print("安装命令:")
print("-" * 70)
print("pip uninstall torch torchvision -y")
print("pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118")
print("-" * 70)
print()

print("【备选方案1】CUDA 12.1 版本（较新，需要新驱动）")
print()
print("安装命令:")
print("-" * 70)
print("pip uninstall torch torchvision -y")
print("pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121")
print("-" * 70)
print()

print("【备选方案2】CUDA 12.4 版本（最新，需要最新驱动）")
print()
print("安装命令:")
print("-" * 70)
print("pip uninstall torch torchvision -y")
print("pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124")
print("-" * 70)
print()

print("=" * 70)
print("安装建议")
print("=" * 70)
print()
print("1. 推荐使用【推荐方案】CUDA 11.8")
print("   - 兼容性最好")
print("   - 无需单独安装 CUDA Toolkit")
print("   - PyTorch 会自动包含所需的 CUDA 运行库")
print()
print("2. 安装后验证:")
print("   python -c \"import torch; print(torch.cuda.is_available())\"")
print("   应该显示: True")
print()
print("3. 如果安装失败，尝试使用国内镜像:")
print("   pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 -i https://pypi.tuna.tsinghua.edu.cn/simple")
print()
print("=" * 70)

